#!/usr/bin/Python
# coding=utf-8
import sys
from thrift.protocol.TBinaryProtocol import TBinaryProtocol
from thrift.transport.TTransport import TMemoryBuffer

sys.path.append('..')

from bdp.i_crawler.i_extractor.ttypes import PageParseInfo
from bdp.i_crawler.i_entity_extractor.ttypes import EntityExtractorInfo
from i_util.normal_proccessor import NormalProccessor
from entity_extractor import EntityExtractor

class EntityExtractorProccessor(NormalProccessor):
    def __init__(self, conf):
        self.log = conf['log']
        self.extract_obj = EntityExtractor(conf)

    def to_string(self, link_info):
        str_entity = None
        try:
            tMemory_b = TMemoryBuffer()
            tBinaryProtocol_b = TBinaryProtocol(tMemory_b)
            link_info.write(tBinaryProtocol_b)
            str_entity = tMemory_b.getvalue()
        except:
            self.log.warning("cann't write EntityExtractorInfo to string")
        return str_entity

    def do_task(self, body):
        parse_info = PageParseInfo()
        try:
            tMemory_o = TMemoryBuffer(body)
            tBinaryProtocol_o = TBinaryProtocol(tMemory_o)
            parse_info.read(tBinaryProtocol_o)
            ret = self.extract_obj.entity_extractor(parse_info)
        except EOFError, e:
            self.log.warning("cann't read PageParseInfo from string")
            return None

        if int(ret.get('CODE', -10000)) < 0:
            return None

        entity_data_list = ret.get('LIST', [])
        msg_list = []
        if len(entity_data_list) > 0:
            for entity_data in entity_data_list:
                if isinstance(entity_data, EntityExtractorInfo):
                    msg_list.append(self.to_string(entity_data))

        else:
            return None
        self.log.info("Topic_id:%s\tsend_msg_num:%s"%(ret.get('TOPIC_ID'), len(msg_list)))
        return msg_list

    def do_output(self, body):
        return True




if __name__ == '__main__':
    pass
